import numpy as np
from sentence_transformers import SentenceTransformer
from gensim.models import Word2Vec
import jieba
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
from tqdm import tqdm
from rapidfuzz import fuzz

class AddressSimilarityTest:
    def __init__(self):
        """初始化模型"""
        self.models = {
            'SBERT': SentenceTransformer('paraphrase-mpnet-base-v2'),
            'BGE': SentenceTransformer('BAAI/bge-base-zh-v1.5'),
            'Word2Vec': None,  # 将在处理数据时训练
            'EditDistance': None,  # 编辑距离算法不需要训练
            'NormalizedDistance': None  # 归一化编辑距离
        }
        
        # 扩充测试用例
        self.test_cases = [
            # 完全相同的地址
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京市海淀区中关村大街1号科技大厦",
             "完全相同"),
            
            # 门牌号差异
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京市海淀区中关村大街2号科技大厦",
             "门牌号不同"),
            ("上海市浦东新区张江高科技园区88号", 
             "上海市浦东新区张江高科技园区86号",
             "门牌号不同"),
            
            # 同街道不同建筑
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京市海淀区中关村大街5号创新广场",
             "同街道不同建筑"),
            ("上海市浦东新区张江高科技园区88号创新大厦", 
             "上海市浦东新区张江高科技园区66号研发中心",
             "同街道不同建筑"),
            
            # 同区不同街道
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京市海淀区学院路15号大学科技园",
             "同区不同街道"),
            ("广州市天河区珠江新城华夏路10号", 
             "广州市天河区体育西路123号",
             "同区不同街道"),
            
            # 同城不同区
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京市朝阳区建国路88号科技中心",
             "同城不同区"),
            ("上海市浦东新区张江高科技园区88号", 
             "上海市徐汇区漕河泾开发区89号",
             "同城不同区"),
            
            # 不同城市
            ("北京市海淀区中关村大街1号科技大厦", 
             "上海市浦东新区张江高科技园区88号",
             "不同城市"),
            ("广州市天河区珠江新城华夏路10号", 
             "深圳市南山区科技园科发路8号",
             "不同城市"),
            
            # 省市区简写
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京海淀中关村大街1号科技大厦",
             "省市区简写"),
            ("广东省广州市天河区珠江新城华夏路10号", 
             "广州天河珠江新城华夏路10号",
             "省市区简写"),
            
            # 简化地址
            ("北京市海淀区中关村大街1号科技大厦", 
             "北京中关村科技大厦",
             "简化地址"),
            ("上海市浦东新区张江高科技园区88号创新大厦", 
             "上海张江创新大厦",
             "简化地址"),
            
            # 别名表达
            ("北京市海淀区中关村大街1号科技大厦A座", 
             "北京中关村科技大厦第一座",
             "别名表达"),
            ("上海市浦东新区张江高科技园区88号创新大厦", 
             "上海张江88号创新中心",
             "别名表达")
        ]

    def get_word2vec_vector(self, address):
        """使用Word2Vec获取地址向量"""
        if self.models['Word2Vec'] is None:
            # 训练Word2Vec模型
            all_addresses = [list(jieba.cut(addr)) for addr, _, _ in self.test_cases]
            all_addresses.extend([list(jieba.cut(addr)) for _, addr, _ in self.test_cases])
            self.models['Word2Vec'] = Word2Vec(
                sentences=all_addresses,
                vector_size=300,
                window=5,
                min_count=1,
                workers=4
            )
        
        # 分词并计算平均词向量
        words = list(jieba.cut(address))
        word_vectors = [self.models['Word2Vec'].wv[word] for word in words if word in self.models['Word2Vec'].wv]
        return np.mean(word_vectors, axis=0) if word_vectors else np.zeros(300)

    def calculate_levenshtein_similarity(self, addr1, addr2, normalized=False):
        """计算编辑距离相似度"""
        if normalized:
            # 使用 rapidfuzz 的归一化比率
            return fuzz.ratio(addr1, addr2) / 100.0
        else:
            # 使用 rapidfuzz 的部分比率
            return fuzz.partial_ratio(addr1, addr2) / 100.0

    def calculate_similarity(self, addr1, addr2, model_name):
        """计算两个地址的相似度"""
        if model_name in ['SBERT', 'BGE']:
            # 使用Transformer模型
            embeddings = self.models[model_name].encode([addr1, addr2])
            return cosine_similarity([embeddings[0]], [embeddings[1]])[0][0]
        elif model_name == 'Word2Vec':
            # 使用Word2Vec
            vec1 = self.get_word2vec_vector(addr1)
            vec2 = self.get_word2vec_vector(addr2)
            return cosine_similarity([vec1], [vec2])[0][0]
        elif model_name == 'EditDistance':
            # 使用编辑距离
            return self.calculate_levenshtein_similarity(addr1, addr2, normalized=False)
        else:  # NormalizedDistance
            # 使用归一化编辑距离
            return self.calculate_levenshtein_similarity(addr1, addr2, normalized=True)

    def run_comparison(self):
        """运行相似度对比测试"""
        results = []
        
        for addr1, addr2, case_type in tqdm(self.test_cases, desc="处理测试用例"):
            case_results = {
                "测试类型": case_type,
                "地址1": addr1,
                "地址2": addr2
            }
            
            # 计算每个模型的相似度
            for model_name in self.models.keys():
                similarity = self.calculate_similarity(addr1, addr2, model_name)
                case_results[f"{model_name}相似度"] = similarity
            
            results.append(case_results)
        
        # 转换为DataFrame并格式化
        df = pd.DataFrame(results)
        
        # 计算评估指标
        similarity_cols = [col for col in df.columns if col.endswith('相似度')]
        metrics = self.calculate_metrics(df, similarity_cols)
        
        # 输出结果
        self.display_results(df, metrics, similarity_cols)
        
        # 保存结果
        df.to_csv('address_similarity_results.csv', index=False)
        metrics.to_csv('similarity_metrics.csv')

    def calculate_metrics(self, df, similarity_cols):
        """计算各种评估指标"""
        metrics_data = {
            '平均相似度': df[similarity_cols].mean(),
            '标准差': df[similarity_cols].std(),
            '最大相似度': df[similarity_cols].max(),
            '最小相似度': df[similarity_cols].min(),
            '完全相同地址相似度': df[df['测试类型'] == '完全相同'][similarity_cols].mean(),
            '不同城市相似度': df[df['测试类型'] == '不同城市'][similarity_cols].mean(),
            '简化地址相似度': df[df['测试类型'] == '简化地址'][similarity_cols].mean()
        }
        
        # 转换为DataFrame并转置，使模型名称成为列名
        metrics_df = pd.DataFrame(metrics_data).round(4)
        return metrics_df

    def display_results(self, df, metrics, similarity_cols):
        """展示结果"""
        print("\n============ 地址相似度测试结果 ============")
        
        # 设置显示格式
        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', None)
        pd.set_option('display.max_colwidth', None)
        pd.set_option('display.precision', 4)
        
        # 按测试类型分组展示结果
        print("\n按类型分组的相似度结果:")
        grouped_results = df.groupby('测试类型')[similarity_cols].mean().round(4)
        print("\n平均相似度（按类型）:")
        print(grouped_results)
        
        # 计算每种类型的标准差
        print("\n相似度标准差（按类型）:")
        grouped_std = df.groupby('测试类型')[similarity_cols].std().round(4)
        print(grouped_std)
        
        # 详细结果
        print("\n详细测试结果:")
        for col in similarity_cols:
            df[col] = df[col].round(4)
        print(df)
        
        # 模型评估指标
        print("\n模型整体评估指标:")
        print(metrics.round(4))
        
        # 模型排名
        print("\n模型综合评分:")
        rankings = self.calculate_rankings(metrics)
        rankings['综合得分'] = rankings['综合得分'].round(4)
        print(rankings)
        
        # 保存结果
        df.to_csv('address_similarity_results.csv', index=False)
        metrics.to_csv('similarity_metrics.csv')
        rankings.to_csv('model_rankings.csv')
        grouped_results.to_csv('similarity_by_type.csv')
        grouped_std.to_csv('std_by_type.csv')
        
        print("\n结果已保存至:")
        print("- address_similarity_results.csv (详细测试结果)")
        print("- similarity_metrics.csv (评估指标)")
        print("- model_rankings.csv (模型排名)")
        print("- similarity_by_type.csv (按类型的平均相似度)")
        print("- std_by_type.csv (按类型的标准差)")

    def calculate_rankings(self, metrics):
        """计算模型综合排名"""
        weights = {
            '平均相似度': 0.15,
            '标准差': -0.15,  # 负权重，因为标准差越小越好
            '完全相同地址相似度': 0.35,
            '不同城市相似度': 0.2,
            '简化地址相似度': 0.15
        }
        
        # 计算加权得分
        scores = {}
        for column in metrics.columns:
            score = 0
            for metric, weight in weights.items():
                if metric in metrics.index:
                    score += metrics.loc[metric, column] * weight
            scores[column] = score
        
        # 创建排名DataFrame
        rankings = pd.DataFrame({
            '综合得分': scores
        }).sort_values('综合得分', ascending=False)
        
        return rankings

if __name__ == "__main__":
    similarity_test = AddressSimilarityTest()
    similarity_test.run_comparison() 